#!/usr/bin/env python3
"""
简化版医学图像处理流程
适用于已经有预测结果文件的情况，直接从normalize开始
"""

import os
import sys
import argparse
import logging
from pathlib import Path
from typing import List

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from auto_pipeline import AutoPipeline

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="简化版医学图像处理流程 - 适用于已有预测结果的情况",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 处理单个受试者（从normalize开始）
  python simple_pipeline.py --subjects CHEN-REN-GENG
  
  # 处理多个受试者
  python simple_pipeline.py --subjects CHEN-REN-GENG DJY BCR
  
  # 模拟运行
  python simple_pipeline.py --subjects CHEN-REN-GENG --dry-run
  
  # 只运行特定步骤
  python simple_pipeline.py --subjects CHEN-REN-GENG --steps normalize segment
  
  # 包含参考文件移动
  python simple_pipeline.py --subjects CHEN-REN-GENG --reference-config reference_config.json
  
  # 指定GPU设备
  python simple_pipeline.py --subjects CHEN-REN-GENG --device gpu:1
  
  # 跳过分割步骤（如果已经有分割结果）
  python simple_pipeline.py --subjects CHEN-REN-GENG --skip-segment
        """
    )
    
    parser.add_argument(
        "--base-dir", 
        type=str, 
        default="..", 
        help="项目根目录路径"
    )
    parser.add_argument(
        "--subjects",
        type=str,
        nargs="+",
        required=True,
        help="要处理的受试者ID列表"
    )
    parser.add_argument(
        "--steps",
        type=str,
        nargs="*",
        choices=["normalize", "segment", "move_reference", "move_masks", "metrics"],
        help="要执行的步骤列表，不指定则执行所有步骤"
    )
    parser.add_argument(
        "--reference-config",
        type=str,
        help="参考文件配置JSON文件路径"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="gpu:0",
        help="分割使用的设备 (gpu:0, gpu:1, cpu等，默认: gpu:0)"
    )
    parser.add_argument(
        "--skip-segment",
        action="store_true",
        help="跳过分割步骤"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="模拟运行，不实际执行操作"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="详细输出"
    )
    
    args = parser.parse_args()
    
    # 设置日志级别
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    # 初始化流程处理器
    pipeline = AutoPipeline(args.base_dir)
    
    # 加载参考文件配置
    reference_configs = {}
    if args.reference_config:
        import json
        try:
            with open(args.reference_config, 'r', encoding='utf-8') as f:
                reference_configs = json.load(f)
        except Exception as e:
            logging.error(f"加载参考文件配置失败: {e}")
    
    # 确定要执行的步骤
    if args.steps:
        selected_steps = args.steps
    else:
        # 默认步骤
        selected_steps = ["normalize", "segment", "move_masks", "move_reference", "metrics"]
        if args.skip_segment:
            selected_steps.remove("segment")
    
    # 步骤映射
    step_mapping = {
        "normalize": ("图像标准化", lambda: pipeline.step3_normalize(args.subjects, args.dry_run)),
        "segment": ("图像分割", lambda: pipeline.step4_segment(args.subjects, args.device, args.dry_run)),
        "move_masks": ("移动分割掩膜", lambda: pipeline.step5_move_seg_masks(args.subjects, args.dry_run)),
        "move_reference": ("移动参考文件", lambda: pipeline.step6_move_reference_files(args.subjects, reference_configs, args.dry_run)),
        "metrics": ("计算评估指标", lambda: pipeline.step7_compute_metrics(args.subjects, args.dry_run))
    }
    
    # 执行选定的步骤
    logging.info("🚀 开始运行简化版医学图像处理流程")
    logging.info(f"受试者列表: {', '.join(args.subjects)}")
    logging.info(f"执行步骤: {', '.join(selected_steps)}")
    logging.info(f"模拟运行: {'是' if args.dry_run else '否'}")
    
    success_count = 0
    for i, step in enumerate(selected_steps, 1):
        step_name, step_func = step_mapping[step]
        logging.info(f"\n🔄 执行步骤 {i}/{len(selected_steps)}: {step_name}")
        
        try:
            if step_func():
                success_count += 1
                logging.info(f"✅ 步骤 {i} 完成: {step_name}")
            else:
                logging.error(f"❌ 步骤 {i} 失败: {step_name}")
        except Exception as e:
            logging.error(f"❌ 步骤 {i} 异常: {step_name} - {e}")
    
    logging.info("=" * 60)
    if success_count == len(selected_steps):
        logging.info("🎉 流程执行成功!")
    else:
        logging.warning(f"⚠️ 流程部分完成: {success_count}/{len(selected_steps)} 个步骤成功")
    logging.info("=" * 60)


if __name__ == "__main__":
    main()